import numpy as np

# # # # # # # # # # # # # # # # # # # # # # #
# # # # # Part 1, interface for RL # # # # #
# 动作编码器模块 - 为强化学习智能体定义动作空间
# # # # # # # # # # # # # # # # # # # # # # #

# 动作字典：定义了强化学习智能体可以执行的所有可能动作
# 每个动作由主命令和子命令组成，格式为 "ActionSet2::主命令;子命令"
# 这些动作涵盖了游戏中智能体的各种行为：警戒、移动、攻击、巡逻、待命等
dictionary_items = [
    'ActionSet2::N/A;N/A',  # 0 - 无动作
    'ActionSet2::Alert;DynamicGuard',  # 1 - 警戒：动态警戒
    'ActionSet2::Alert;Crouch',  # 2 - 警戒：蹲伏
    'ActionSet2::Alert;Prone',  # 3 - 警戒：俯卧
    'ActionSet2::Alert;StaticAlert',  # 4 - 警戒：静态警戒
    'ActionSet2::Alert;AggressivePersue',  # 5 - 警戒：激进追击
    'ActionSet2::Moving;+X',  # 6 - 移动：沿X轴正方向
    'ActionSet2::Moving;+Y',  # 7 - 移动：沿Y轴正方向
    'ActionSet2::Moving;-X',  # 8 - 移动：沿X轴负方向
    'ActionSet2::Moving;-Y',  # 9 - 移动：沿Y轴负方向
    'ActionSet2::NormalAttacking;T1-0',  # 10 - 普通攻击：攻击队伍1的第0个目标
    'ActionSet2::NormalAttacking;T1-1',  # 11 - 普通攻击：攻击队伍1的第1个目标
    'ActionSet2::NormalAttacking;T1-2',  # 12 - 普通攻击：攻击队伍1的第2个目标
    'ActionSet2::NormalAttacking;T1-3',  # 13 - 普通攻击：攻击队伍1的第3个目标
    'ActionSet2::NormalAttacking;T1-4',  # 14 - 普通攻击：攻击队伍1的第4个目标

    'ActionSet2::NormalAttacking;T0-0',  # 15 - 普通攻击：攻击队伍0的第0个目标
    'ActionSet2::NormalAttacking;T0-1',  # 16 - 普通攻击：攻击队伍0的第1个目标
    'ActionSet2::NormalAttacking;T0-2',  # 17 - 普通攻击：攻击队伍0的第2个目标
    'ActionSet2::NormalAttacking;T0-3',  # 18 - 普通攻击：攻击队伍0的第3个目标
    'ActionSet2::NormalAttacking;T0-4',  # 19 - 普通攻击：攻击队伍0的第4个目标

    'ActionSet2::Patrol;+X',  # 20 - 巡逻：沿X轴正方向
    'ActionSet2::Patrol;+Y',  # 21 - 巡逻：沿Y轴正方向
    'ActionSet2::Patrol;-X',  # 22 - 巡逻：沿X轴负方向
    'ActionSet2::Patrol;-Y',  # 23 - 巡逻：沿Y轴负方向

    'ActionSet2::Idle;AsFarAsPossible',  # 24 - 待命：尽可能远离目标
    'ActionSet2::Idle;StayWhenTargetInRange',  # 25 - 待命：目标在范围内时停留
    'ActionSet2::Idle;StayWhenTargetInHalfRange',  # 26 - 待命：目标在半程范围内时停留
]

# 动作总数：智能体可以选择的动作数量
dictionary_n_actions = len(dictionary_items)

# 数字到动作的映射字典：将动作索引转换为动作字符串
digit2act_dictionary = {
    i: dictionary_items[i] for i, item in enumerate(dictionary_items)
}

# 动作到数字的映射字典：将动作字符串转换为动作索引
act2digit_dictionary = {
    dictionary_items[i]: i for i, item in enumerate(dictionary_items)
}

# # # # # # # # # # # # # # # # # # # # # # #
# # # # # Part 2, translate actions # # # # #
# 智能体属性映射：将JSON格式的智能体数据转换为本地属性
# # # # # # # # # # # # # # # # # # # # # # #

# JSON字段到本地属性的映射列表
# 将从游戏引擎接收的JSON数据转换为智能体对象的本地属性
agent_json2local_attrs = [
    # json键名          ->    本地属性名
    ('agentAlive', 'alive'),           # 智能体存活状态
    ('agentTeam', 'team'),             # 智能体所属队伍
    ('indexInTeam', 'index'),          # 在队伍中的索引
    ('uId', 'uid_remote'),             # 远程唯一标识符
    ('maxMoveSpeed', 'max_speed'),     # 最大移动速度
    ('agentLocationArr', 'location'),  # 位置坐标数组
    ('agentRotationArr', 'rotation'),  # 旋转角度数组
    ('agentScaleArr', 'scale3'),       # 缩放比例数组
    ('agentVelocityArr', 'velocity'),  # 速度向量数组
    ('agentHp', 'hp'),                 # 生命值
    ('weaponCD', 'weapon_cd'),         # 武器冷却时间
    ('type', 'type'),                  # 智能体类型
]

# 动作编码函数：将动作参数转换为数字数组格式
# 这个函数将人类可读的动作参数转换为神经网络可以处理的数字格式
def encode_action_as_digits(main_cmd, sub_cmd, x=None, y=None, z=None, UID=None, T=None, T_index=None):
    """
    将动作编码为数字数组，便于神经网络处理

    Args:
        main_cmd: 主命令（如 "Moving", "Attacking", "Idle" 等）
        sub_cmd: 子命令（如 "+X", "+Y", "T1-0" 等）
        x, y, z: 坐标参数（可选）
        UID: 目标唯一标识符（可选）
        T: 目标队伍编号（可选）
        T_index: 目标在队伍中的索引（可选）

    Returns:
        包含8个元素的numpy数组，表示编码后的动作
    """
    # 主命令编码器：将文本命令映射为数字
    main_cmd_encoder = {
        "Idle": 0,          # 待命
        "Guard": 1,         # 警戒
        "Moving": 2,        # 移动
        "GrenadeAttacking": 3,  # 手榴弹攻击
        "NormalAttacking": 4,    # 普通攻击
    }

    # 子命令编码器：将方向和攻击目标映射为数字
    sub_cmd_encoder = {
        "N/A": 0,   # 无动作/不适用
        '+X': 1,    # X轴正方向
        '+X+Y': 2,  # X轴正方向+Y轴正方向
        '+Y': 3,    # Y轴正方向
        '-X+Y': 4,  # X轴负方向+Y轴正方向
        '-X': 5,    # X轴负方向
        '-X-Y': 6,  # X轴负方向+Y轴负方向
        '-Y': 7,    # Y轴负方向
        '+X-Y': 8,  # X轴正方向+Y轴负方向

        # 包含Z轴的3D移动方向
        '+X+Z': 9,      # X轴正方向+Z轴正方向
        '+Y+Z': 10,     # Y轴正方向+Z轴正方向
        '-X+Z': 11,     # X轴负方向+Z轴正方向
        '-Y+Z': 12,     # Y轴负方向+Z轴正方向
        '+X+Y+Z': 13,   # X轴正方向+Y轴正方向+Z轴正方向
        '-X+Y+Z': 14,   # X轴负方向+Y轴正方向+Z轴正方向
        '-X-Y+Z': 15,   # X轴负方向+Y轴负方向+Z轴正方向
        '+X-Y+Z': 16,   # X轴正方向+Y轴负方向+Z轴正方向

        # 负Z轴方向的3D移动
        '+X-Z': 17,     # X轴正方向+Z轴负方向
        '+Y-Z': 18,     # Y轴正方向+Z轴负方向
        '-X-Z': 19,     # X轴负方向+Z轴负方向
        '-Y-Z': 20,     # Y轴负方向+Z轴负方向
        '+X+Y-Z': 21,   # X轴正方向+Y轴正方向+Z轴负方向
        '-X+Y-Z': 22,   # X轴负方向+Y轴正方向+Z轴负方向
        '-X-Y-Z': 23,   # X轴负方向+Y轴负方向+Z轴负方向
        '+X-Y-Z': 24,   # X轴正方向+Y轴负方向+Z轴负方向

        # 纯Z轴移动
        '+Z': 25,       # Z轴正方向
        '-Z': 26        # Z轴负方向
    }

    # 调试输出：当动作为待命时打印信息
    if main_cmd == 0:
        print('Idle')

    # 返回编码后的动作数组
    # 数组格式：[主命令编码, 子命令编码, x坐标, y坐标, z坐标, 目标UID, 目标队伍, 目标索引]
    # 未指定的参数设为无穷大(np.inf)表示无效
    return np.array([
        main_cmd_encoder[main_cmd],
        sub_cmd_encoder[sub_cmd],
        x if x is not None else np.inf,
        y if y is not None else np.inf,
        z if z is not None else np.inf,
        UID if UID is not None else np.inf,
        T if T is not None else np.inf,
        T_index if T_index is not None else np.inf
    ])


# 动作解码函数：将数字数组格式转换为人类可读的动作字符串
def decode_action_as_string(digits):
    """
    将数字数组解码为动作字符串，便于人类理解和调试

    Args:
        digits: 包含8个元素的数字数组，表示编码后的动作

    Returns:
        人类可读的动作字符串
    """
    # 主命令解码器：将数字映射回文本命令
    main_cmd_decoder = {
        0:"Idle",                 # 待命
        1:"Guard",                # 警戒
        2:"Moving",               # 移动
        3:"GrenadeAttacking",     # 手榴弹攻击
        4:"NormalAttacking",      # 普通攻击
    }

    # 子命令解码器：将数字映射回方向和攻击目标
    sub_cmd_decoder = {
        0:"N/A",      # 无动作/不适用
        1:'+X',       # X轴正方向
        2:'+X+Y',     # X轴正方向+Y轴正方向
        3:'+Y',       # Y轴正方向
        4:'-X+Y',     # X轴负方向+Y轴正方向
        5:'-X',       # X轴负方向
        6:'-X-Y',     # X轴负方向+Y轴负方向
        7:'-Y',       # Y轴负方向
        8:'+X-Y',     # X轴正方向+Y轴负方向

        # 包含Z轴的3D移动方向
        9: '+X+Z',      # X轴正方向+Z轴正方向
        10:'+Y+Z',      # Y轴正方向+Z轴正方向
        11:'-X+Z',      # X轴负方向+Z轴正方向
        12:'-Y+Z',      # Y轴负方向+Z轴正方向
        13:'+X+Y+Z',    # X轴正方向+Y轴正方向+Z轴正方向
        14:'-X+Y+Z',    # X轴负方向+Y轴正方向+Z轴正方向
        15:'-X-Y+Z',    # X轴负方向+Y轴负方向+Z轴正方向
        16:'+X-Y+Z',    # X轴正方向+Y轴负方向+Z轴正方向

        # 负Z轴方向的3D移动
        17:'+X-Z',      # X轴正方向+Z轴负方向
        18:'+Y-Z',      # Y轴正方向+Z轴负方向
        19:'-X-Z',      # X轴负方向+Z轴负方向
        20:'-Y-Z',      # Y轴负方向+Z轴负方向
        21:'+X+Y-Z',    # X轴正方向+Y轴正方向+Z轴负方向
        22:'-X+Y-Z',    # X轴负方向+Y轴正方向+Z轴负方向
        23:'-X-Y-Z',    # X轴负方向+Y轴负方向+Z轴负方向
        24:'+X-Y-Z',    # X轴正方向+Y轴负方向+Z轴负方向

        # 纯Z轴移动
        25:'+Z',        # Z轴正方向
        26:'-Z'         # Z轴负方向
    }

    # 从数字数组中解码各个参数
    main_cmd = main_cmd_decoder[digits[0]]
    sub_cmd = sub_cmd_decoder[digits[1]]
    x = digits[2] if np.isfinite(digits[2]) else None
    y = digits[3] if np.isfinite(digits[3]) else None
    z = digits[4] if np.isfinite(digits[4]) else None
    UID = digits[5] if np.isfinite(digits[5]) else None
    T = digits[6] if np.isfinite(digits[6]) else None
    T_index = digits[7] if np.isfinite(digits[7]) else None

    # 根据主命令类型构建动作字符串
    if main_cmd == "Alert":
        # 警戒类动作
        res = 'ActionSet2::Alert;%s' % sub_cmd
        assert res in dictionary_items, '指令错误无法解析'

    elif main_cmd == "Moving":
        # 移动类动作
        if sub_cmd == 'N/A' and x is not None and y is not None and z is not None:
            # 直接坐标移动
            res = 'ActionSet2::Moving;X=%f Y=%f Z=%f' % (x, y, z)
        elif sub_cmd != 'N/A':
            # 方向性移动
            res = 'ActionSet2::Moving;%s' % sub_cmd

    elif main_cmd == "GrenadeAttacking":
        # 手榴弹攻击类动作
        if sub_cmd == 'N/A':
            # 直接坐标攻击
            res = 'ActionSet2::GrenadeAttacking;X=%f Y=%f Z=%f' % (x, y, z)
        else:
            # 方向性攻击
            res = 'ActionSet2::GrenadeAttacking;%s' % sub_cmd

    elif main_cmd == "Patrol":
        # 巡逻类动作
        if sub_cmd == 'N/A':
            # 直接坐标巡逻
            res = 'ActionSet2::Patrol;X=%f Y=%f Z=%f' % (x, y, z)
        else:
            # 方向性巡逻
            res = 'ActionSet2::Patrol;%s' % sub_cmd

    elif main_cmd == "NormalAttacking":
        # 普通攻击类动作
        assert sub_cmd == 'N/A', '指令错误无法解析'
        if UID is not None:
            # 攻击特定UID的目标
            res = 'ActionSet2::NormalAttacking;UID-%d' % UID
        else:
            # 攻击特定队伍的特定目标
            res = 'ActionSet2::NormalAttacking;%d;%d' % (T, T_index)

    elif main_cmd == "N/A":
        # 无动作
        res = 'ActionSet2::N/A;N/A'

    elif main_cmd == "Idle":
        # 待命状态
        res = 'ActionSet2::Idle;N/A'

    elif main_cmd == "Guard":
        # 警戒类动作
        if sub_cmd == 'N/A' and x is not None and y is not None and z is not None:
            # 直接坐标警戒
            res = 'ActionSet2::Guard;X=%f Y=%f Z=%f' % (x, y, z)
        elif sub_cmd != 'N/A':
            # 方向性警戒
            res = 'ActionSet2::Guard;%s' % sub_cmd
    else:
        print('指令错误无法解析')
        assert False
    return res


# # # # # # # # # # # # # # # # # # # # # # #
# # # # # Part 3, agent init defaults # # # #
# 智能体默认属性配置
# # # # # # # # # # # # # # # # # # # # # # #

# 智能体属性默认值字典
# 定义了新创建智能体时的所有默认属性
AgentPropertyDefaults = {
    'ClassName': 'RLA_CAR',              # 类名：强化学习智能体载具
    'DebugAgent': False,                 # 调试模式开关
    'AgentTeam': 0,                      # 所属队伍（默认队伍0）
    'IndexInTeam': 0,                    # 在队伍中的索引
    'UID': 0,                            # 唯一标识符
    'MaxMoveSpeed': 600,                 # 最大移动速度（单位：游戏单位/秒）
    'InitLocation': {'x': 0, 'y': 0, 'z': 0, },  # 初始位置坐标
    'InitRotation': {'x': 0, 'y': 0, 'z': 0, },  # 初始旋转角度
    'AgentScale': {'x': 1, 'y': 1, 'z': 1, },    # 智能体缩放比例
    'InitVelocity': {'x': 0, 'y': 0, 'z': 0, },  # 初始速度向量
    'AgentHp': 100,                      # 初始生命值
    "WeaponCD": 1,                       # 武器冷却时间（秒）
    "IsTeamReward": True,                # 是否使用团队奖励
    "Type": "",                          # 智能体类型（字符串）
    "DodgeProb": 0.8,                    # 闪避概率（0.0-1.0）
    "ExplodeDmg": 25,                    # 爆炸伤害值
    "FireRange": 1000.0,                 # 武器射程（单位：游戏单位）
    "GuardRange": 1400.0,                # 警戒范围（单位：游戏单位）
    "PerceptionRange": 1500.0,           # 感知范围（单位：游戏单位）
    'Color': '(R=0,G=1,B=0,A=1)',       # 智能体颜色（RGBA格式）
    "FireRange": 1000,                   # 重复定义的射程（可能是历史遗留）
    'RSVD1': '',                         # 保留字段1
    'RSVD2': '',                         # 保留字段2
}


# # # # # # # # # # # # # # # # # # # # # # #
# # # # # Part 4, utility functions # # # #
# 工具函数：二进制兼容性检查和字符串编码
# # # # # # # # # # # # # # # # # # # # # # #

# 检查数字是否可以被浮点数精确表示
# 这是一个用于避免"蝴蝶效应"问题的函数，确保数值精度不会导致意外的行为
def binary_friendly(x):
    """
    检查数字x是否可以被float16精确表示，避免浮点数精度问题

    Args:
        x: 要检查的数字

    Returns:
        bool: 如果可以被精确表示则返回True，否则返回False
    """
    # 将数字转换为float16和float64格式进行比较
    y_f16 = np.array(x, dtype=np.float16)
    y_f64 = np.array(x, dtype=np.float64)
    t = y_f64 - y_f16
    assert t.dtype == np.float64
    return (t == 0)

# 动作数字编码的长度限制
ActDigitLen = 100

# 将动作字符串转换为数字数组
def strActionToDigits(act_string):
    """
    将动作字符串转换为数字数组，便于神经网络处理

    Args:
        act_string: 动作字符串

    Returns:
        长度为ActDigitLen的数字数组
    """
    # 将字符串的每个字符转换为ASCII码
    t = [ord(c) for c in act_string]
    d_len = len(t)
    # 检查字符串长度是否超出限制
    assert d_len <= ActDigitLen, ("Action string is tooo long! Don't be wordy. Or you can increase ActDigitLen above.")
    # 用-1填充剩余位置
    pad = [-1 for _ in range(ActDigitLen - d_len)]
    return (t + pad)

# 将数字数组转换为动作字符串
def digitsToStrAction(digits):
    """
    将数字数组转换回动作字符串，用于解码

    Args:
        digits: 数字数组

    Returns:
        动作字符串
    """
    # 如果所有元素都是0，返回默认无动作
    if all([a == 0 for a in digits]):
        return 'ActionSet3::N/A;N/A'
    # 将正数字转换为对应的ASCII字符，忽略负数
    arr = [chr(d) for d in digits.astype(int) if d >= 0]
    return ''.join(arr)


# # # # # # # # # # # # # # # # # # # # # # #
# # # # # Part 5, legacy code (commented) # # # #
# 历史代码（已注释）- 单数字编码方案，当前未使用
# # # # # # # # # # # # # # # # # # # # # # #

"""
以下是被注释掉的单数字编码代码，用于将3D坐标和动作类型编码为单个数字。
这个方案目前没有被使用，但保留用于可能的未来实现。

核心思想：
- 将地图划分为网格
- 将3D坐标转换为网格索引
- 将动作类型和坐标信息组合成一个数字
- 使用混合进制编码来节省空间

示例：
h_map_center = (-7290.0, 6010.0)  # 地图中心坐标
h_grid_size = 400                 # 水平网格大小
v_ground = 340                    # 垂直基准高度
v_grid_size = 1000                # 垂直网格大小

x_arr = np.array([h_map_center[0]+v_grid_size*i for i in range(-20, 20)])   # X轴坐标数组，40个值
y_arr = np.array([h_map_center[1]+v_grid_size*i for i in range(-20, 20)])   # Y轴坐标数组，40个值
z_arr = np.array([v_ground+v_grid_size*i for i in range(4)])                # Z轴坐标数组，4个值

def _2digit(main_cmd, x, y, z):
    # 将3D坐标和动作类型编码为单个数字
    z_logit = np.argmin(np.abs(z - z_arr))     # 找到最近的Z轴索引
    x_logit = np.argmin(np.abs(x - x_arr))     # 找到最近的X轴索引
    y_logit = np.argmin(np.abs(y - y_arr))     # 找到最近的Y轴索引
    if main_cmd=='SpecificMoving': cmd_logit = 0
    elif main_cmd=='PatrolMoving': cmd_logit = 1
    ls_mod = [1,40,1600,6400]                  # 混合适制基数
    offset = 1000                               # 偏移量
    x = np.array([x_logit, y_logit, z_logit, cmd_logit])
    y = np.dot(x, ls_mod)+offset
    return y

def _2coordinate(x):
    # 将单个数字解码回3D坐标和动作类型
    offset = 1000
    ls_mod = [1,40,1600,6400]
    x = x - offset
    res = []
    for mod in reversed(ls_mod):
        tmp = x // mod
        x = x - tmp*mod
        res.append(tmp)
    res = list(reversed(res))
    x_logit, y_logit, z_logit, cmd_logit = res
    if    cmd_logit == 0 : main_cmd ='SpecificMoving'
    elif  cmd_logit == 1 : main_cmd ='PatrolMoving'
    x = x_arr[x_logit]
    y = y_arr[y_logit]
    z = z_arr[z_logit]
    return main_cmd, x, y, z
"""